# -*- coding=utf-8 -*-
# 2021.
# @zlk 
# zhanglk960127@163.com

# 寻找transpose的ortholog

## 2023.11.19 并且候选的两个转座基因的100个基因内不能有共线性基因
import sys
import os
import argparse
import re
import math
import multiprocessing
parser = argparse.ArgumentParser(description='pangenome')
parser.add_argument('-m', type=str, help='output folder of mSynOrths')
args = parser.parse_args()
msyn_fold=args.m
# msyn_fold,frame_fold,out_file=args.m,args.f,args.o

def get_ortho_func(msyn_fold,msyn,blast_file,syn_file,total_gene_length,gene_order2,gene_order1):
    op_syn_file=open(msyn_fold+'/'+msyn+'/'+syn_file,'r')
    first_list,second_list=[],[]
    syn_dict={}
    for line in op_syn_file:
        if line.startswith('##'):
            continue
        line_list=line.strip().split('\t')
        if line_list[0] not in syn_dict.keys():
            syn_dict[line_list[0]]=[]
        syn_dict[line_list[0]].append(line_list[1])
        first_list.append(line_list[1])
        second_list.append(line_list[0])
    op_syn_file.close()
    
    op_blast_file=open(msyn_fold+'/'+msyn+'/'+blast_file,'r')
    besthit_dict={}
    besthit_dict2={}
    last=''
    have_list=[]
    for line in op_blast_file:
        line_list=line.strip().split('\t')
        if line_list[0]!=last:
            last=line_list[0]
            if line_list[0] not in first_list and  line_list[1] not in second_list:
                if float(line_list[2])>75 and float(line_list[-2])<10**-20 and float(line_list[3])/total_gene_length[line_list[0]]>0.6 and float(line_list[3])/total_gene_length[line_list[1]]>0.6:
                    if line_list[0] not in besthit_dict.keys() and line_list[0] not in have_list:
                        besthit_dict[line_list[0]]=[line_list[1],line_list[2],line_list[-2]]
            else: 
                have_list.append(line_list[0])
        if line_list[1] not in besthit_dict2.keys():
            besthit_dict2[line_list[1]]=[line_list[0],line_list[2],line_list[-2]]
        else:
            if float(line_list[2])>float(besthit_dict2[line_list[1]][1]):
                besthit_dict2[line_list[1]]=[line_list[0],line_list[2],line_list[-2]]
                
                
    op_blast_file.close()
    out_ortholog_file=open(msyn_fold+'/'+msyn+'/'+blast_file[:-5]+'ortho','w')
    have_list=[]
    # print(len(besthit_dict))
    for key,value in besthit_dict.items():
        if key in have_list or value[0] in have_list:
            continue
        try:
            if key==besthit_dict2[value[0]][0] and float(besthit_dict2[value[0]][1])>75 and float(besthit_dict2[value[0]][2])<10**-20:
                ### 检验两个基因的临近有没有共线性基因
                gene_index1=gene_order1.index(key)
                gene_index2=gene_order2.index(value[0])
                gene_list1=gene_order1[gene_index1-50:gene_index1+50]
                syn_list=[]
                for i in gene_order2[gene_index2-50:gene_index2+50]:
                    if i in syn_dict.keys():
                        syn_list+=syn_dict[i]
                intersection = set(gene_list1) & set(syn_list)
                if len(intersection)==0:
                    out_ortholog_file.write(value[0]+'\t'+key+'\t'+value[1]+'\t'+value[2]+'\n')
                    have_list.append(key)
                    have_list.append(value[0])
        except KeyError:
            continue
    out_ortholog_file.close()
    
total_tandem_dict={}
total_proximal_dict={}
total_syn_list=[]
# op_syn_file=open(msyn_fold+'/Total_species_syntenic_gene_pairs.txt','r')
# for line in op_syn_file:
#     line_list=line.strip().split('\t')
#     for i_list in line_list:
#         num=int(re.findall('\d+',i_list)[0])
#         for i in i_list[num+1:].split(','):
#             if i!='':
#                 total_syn_list.append(i)
total_gene_length={}
total_gene_order={}
for msyn in os.listdir(msyn_fold):
    if msyn.startswith('mSynF'): 
        op_pos_file=open(msyn_fold+'/'+msyn+'/species.gff_pos')
        file_num=re.findall('\d+',msyn)[0]
        total_gene_order[file_num]=[]
        for line in op_pos_file:
            line_list=line.strip().split()
            total_gene_length[line_list[0]]=int(line_list[2])/3
            total_gene_order[file_num].append(line_list[0])
                    
pool = multiprocessing.Pool(40)             
for msyn in os.listdir(msyn_fold):
    if msyn.startswith('mSynF'):
        for blast_file in os.listdir(msyn_fold+'/'+msyn):
            if blast_file.endswith('.blast'):
                syn_file=blast_file[:-5]+'syn'
                ortho_file=msyn_fold+'/'+msyn+'/'+blast_file[:-5]+'ortho'
                # if ortho_file in os.listdir(msyn_fold+'/'+msyn):
                #     continue
                num1=re.findall('\d+',blast_file)[0]
                num2=re.findall('\d+',blast_file)[1]
                pool.apply_async(get_ortho_func,(msyn_fold,msyn,blast_file,syn_file,total_gene_length,total_gene_order[num1],total_gene_order[num2]))
                # get_ortho_func(msyn_fold,msyn,blast_file,syn_file,total_gene_length,total_gene_order[num1],total_gene_order[num2])
pool.close()
pool.join()
                        
                        
                
            
        
        
        
        
